# query_api.py
import os
from fastapi import FastAPI, HTTPException, Form
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from dotenv import load_dotenv
from query import search_similar_chunks, generate_answer
from pymilvus import connections, Collection
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 加载环境变量
load_dotenv()

# 初始化FastAPI应用
app = FastAPI(title="文档查询API", description="基于Milvus和大模型的文档查询服务")

# 配置Milvus连接
@app.on_event("startup")
async def startup_event():
    try:
        connections.connect(
            host=os.getenv('MILVUS_HOST', 'localhost'),
            port=os.getenv('MILVUS_PORT', '19530')
        )
        logger.info(f"成功连接到Milvus服务器")
        # 加载集合
        global collection
        collection_name = os.getenv('MILVUS_COLLECTION_NAME', 'bge_docs')
        collection = Collection(collection_name)
        collection.load()
        logger.info(f"成功加载Milvus集合: {collection_name}")
    except Exception as e:
        logger.error(f"初始化Milvus连接时出错: {str(e)}")
        raise HTTPException(status_code=500, detail=f"初始化Milvus连接时出错: {str(e)}")

# 请求模型
class QueryRequest(BaseModel):
    question: str
    top_k: int = 3

# 响应模型
class QueryResponse(BaseModel):
    answer: str
    sources: list
    status: str

# POST查询接口 - 同时支持带斜杠和不带斜杠的路径，支持表单数据
@app.post("/query", response_model=QueryResponse)
@app.post("/query/", response_model=QueryResponse)
async def query_documents(question: str = Form(...), top_k: int = Form(3)):
    """
    文档查询接口
    - 接收用户问题和返回结果数量
    - 在Milvus中搜索相似文本块
    - 调用大模型生成答案
    - 返回答案和参考来源
    """
    try:
        if not question or not question.strip():
            raise HTTPException(status_code=400, detail="问题不能为空")

        logger.info(f"接收查询请求: {question}, top_k: {top_k}")

        # 搜索相关文本块
        similar_chunks = search_similar_chunks(question, top_k=top_k)

        # 生成答案
        answer = generate_answer(question, similar_chunks)

        # 准备参考来源
        sources = []
        if similar_chunks:
            sources = [{
                "file_name": chunk['file_name'],
                "relevance": round(chunk['score'], 3)
            } for chunk in similar_chunks]

        logger.info(f"查询完成，返回结果")

        return JSONResponse(content={
            "answer": answer,
            "sources": sources,
            "status": "success"
        })
    except Exception as e:
        logger.error(f"查询时出错: {str(e)}")
        raise HTTPException(status_code=500, detail=f"查询时出错: {str(e)}")

# 健康检查接口
@app.get("/")
async def health_check():
    return {
        "status": "healthy",
        "message": "文档查询API服务正常运行，请使用POST /query接口进行查询"
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)